# # Copyright (c) 2019-present, HuggingFace Inc.

import os
import logging
import random
from itertools import chain
from argparse import ArgumentParser
from pprint import pformat

import torch
import torch.nn.functional as F

from transformers import OpenAIGPTLMHeadModel, GPT2LMHeadModel, BertTokenizer

import run_pplm

from flask import Flask, render_template, request,jsonify
import execute
import time
import threading
import jieba
from flask_cors import *
import json

"""
定义心跳检测函数
"""
def heartbeat():
    print (time.strftime('%Y-%m-%d %H:%M:%S - heartbeat', time.localtime(time.time())))
    timer = threading.Timer(60, heartbeat)
    timer.start()
timer = threading.Timer(60, heartbeat)
timer.start()

# -*- coding: utf-8 -*-

import requests
import random
import json
from hashlib import md5

# Set your own appid/appkey.
appid = ''
appkey = ''

# For list of language codes, please refer to `https://api.fanyi.baidu.com/doc/21`
from_lang = 'zh'
to_lang =  'en'

endpoint = 'http://api.fanyi.baidu.com'
path = '/api/trans/vip/translate'
url = endpoint + path

def make_md5(s, encoding='utf-8'):
               return md5(s.encode(encoding)).hexdigest()

SPECIAL_TOKENS = ["[CLS]", "[SEP]", "[PAD]", "[speaker1]", "[speaker2]"]


def top_filtering(logits, top_k=0, top_p=0.0, threshold=-float('Inf'), filter_value=-float('Inf')):
    """ Filter a distribution of logits using top-k, top-p (nucleus) and/or threshold filtering
        Args:
            logits: logits distribution shape (vocabulary size)
            top_k: <=0: no filtering, >0: keep only top k tokens with highest probability.
            top_p: <=0.0: no filtering, >0.0: keep only a subset S of candidates, where S is the smallest subset
                whose total probability mass is greater than or equal to the threshold top_p.
                In practice, we select the highest probability tokens whose cumulative probability mass exceeds
                the threshold top_p.
            threshold: a minimal threshold to keep logits
    """
    assert logits.dim() == 1  # Only work for batch size 1 for now - could update but it would obfuscate a bit the code
    top_k = min(top_k, logits.size(-1))
    if top_k > 0:
        # Remove all tokens with a probability less than the last token in the top-k tokens
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    if top_p > 0.0:
        # Compute cumulative probabilities of sorted tokens
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probabilities = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens with cumulative probability above the threshold
        sorted_indices_to_remove = cumulative_probabilities > top_p
        # Shift the indices to the right to keep also the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # Back to unsorted indices and set them to -infinity
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[indices_to_remove] = filter_value

    indices_to_remove = logits < threshold
    logits[indices_to_remove] = filter_value

    return logits


def build_input_from_segments(history, reply, tokenizer, with_eos=True):
    """ Build a sequence of input from 3 segments: persona, history and last reply """
    bos, eos, pad, speaker1, speaker2 = tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS)
    sequence = [[bos]] + history + [reply + ([eos] if with_eos else [])]
    sequence = [sequence[0]] + [[speaker2 if i % 2 else speaker1] + s for i, s in enumerate(sequence[1:])]
    instance = {}
    instance["input_ids"] = list(chain(*sequence))
    instance["token_type_ids"] = [bos] + [speaker2 if i % 2 else speaker1 for i, s in enumerate(sequence[1:])
                                          for _ in s]
    return instance, sequence


def sample_sequence(history, tokenizer, model, args, current_output=None):
    special_tokens_ids = tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS)
    if current_output is None:
        current_output = []

    for i in range(args.max_length):
        instance, sequence = build_input_from_segments(history, current_output, tokenizer, with_eos=False)
        input_ids = torch.tensor(instance["input_ids"], dtype=torch.long, device=args.device).unsqueeze(0)
        token_type_ids = torch.tensor(instance["token_type_ids"], dtype=torch.long, device=args.device).unsqueeze(0)
        #print(model(input_ids, token_type_ids=token_type_ids))
        logits, *_ = model(input_ids, token_type_ids=token_type_ids)
        #print(logits)
        logits = logits[0, -1, :] / args.temperature
        #logits = logits[0 : -1] / args.temperature
        logits = top_filtering(logits, top_k=args.top_k, top_p=args.top_p)
        probs = F.softmax(logits, dim=-1)

        prev = torch.topk(probs, 1)[1] if args.no_sample else torch.multinomial(probs, 1)
        if i < args.min_length and prev.item() in special_tokens_ids:
            while prev.item() in special_tokens_ids:
                prev = torch.multinomial(probs, num_samples=1)

        if prev.item() in special_tokens_ids:
            break
        current_output.append(prev.item())

    return current_output




app = Flask(__name__,static_url_path="/static")
CORS(app, supports_credentials=True)  # 设置跨域

history = []

@app.route('/message', methods=['POST'])
@cross_origin()
#"""定义应答函数，用于获取输入信息并返回相应的答案"""
def run1():
    # 从请求中获取参数信息
    ##php
    
    global history

    random_reply = ["hh,你接着说","这...","我想想","开心点","有点难，我之后查查相关资料"]
    blacktokens = []
    with open("blacktokens.txt","r",encoding="utf-8") as f:
        for i in f:
            blacktokens.append(i.strip())
    print(blacktokens)


    data = json.loads(request.form['data'])
    req_msg_req = data['msg']
    ##python
    # req_msg_req = request.form['msg']
    # 将语句使用结巴分词进行分词
    # req_msg = " ".join(jieba.cut(req_msg_req))
    # print(req_msg)
    raw_text = req_msg_req
    while not raw_text:
        print('Prompt should not be empty!')
    raw_text = " ".join(list(raw_text.replace(" ", "")))
    history.append(tokenize(raw_text))
    with torch.no_grad():
        out_ids = sample_sequence(history, tokenizer, model, args)
    history.append(out_ids)
    history = history[-(2 * args.max_history + 1):]
    out_text = tokenizer.decode(out_ids, skip_special_tokens=True)
    print(out_text)

    out = ''.join(out_text.split())
    #print(out_text)

    out = out.replace("付费","找专人")
    
    #print(out)
    for i in blacktokens:
        if out.find(i) > 0:
            print(random.choice(random_reply))
            break
    else:
        print(out)

    # query = out_text
    #
    # # Generate salt and sign
    # salt = random.randint(32768, 65536)
    # sign = make_md5(appid + query + str(salt) + appkey)
    #
    # # Build request
    # headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    # payload = {'appid': appid, 'q': query, 'from': "zh", 'to': "en", 'salt': salt, 'sign': sign}
    #
    # # Send request
    # r = requests.post(url, params=payload, headers=headers)
    # result = r.json()
    # rst = eval(json.dumps(result, indent=4, ensure_ascii=False))["trans_result"][0]["dst"]
    # print(rst)
    #
    # args1 = run_pplm.extract_parsm()
    # convc = rst
    # args1.cond_text = convc
    # print(args1)
    # res = run_pplm.run_pplm_example(**vars(args1))
    # print(res)
    #
    # query = res[0].replace("\n", "")
    # print(query)
    #
    # # Generate salt and sign
    # salt = random.randint(32768, 65536)
    # sign = make_md5(appid + query + str(salt) + appkey)
    #
    # # Build request
    # headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    # payload = {'appid': appid, 'q': query, 'from': "en", 'to': "zh", 'salt': salt, 'sign': sign}
    #
    # # Send request
    # r = requests.post(url, params=payload, headers=headers)
    # result = r.json()
    # rst = eval(json.dumps(result, indent=4, ensure_ascii=False))["trans_result"][0]["dst"]
    # print(rst)
    #
    res_msg = out
    # 如果接受到的内容为空，则给出相应的回复
    if res_msg == ' ':
        res_msg = '请与我聊聊天吧'

    return jsonify({'text': res_msg})

@app.route("/")
def index():
    return render_template("index.html")

def run():
    parser = ArgumentParser()
    parser.add_argument('--gpt2', action='store_true', help="use gpt2")
    parser.add_argument("--model_checkpoint", type=str, default="", help="Path, url or short name of the model")
    parser.add_argument("--max_history", type=int, default=5, help="Number of previous utterances to keep in history")
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
                        help="Device (cuda or cpu)")

    parser.add_argument("--no_sample", action='store_true', help="Set to use greedy decoding instead of sampling")
    parser.add_argument("--max_length", type=int, default=30, help="Maximum length of the output utterances")
    parser.add_argument("--min_length", type=int, default=1, help="Minimum length of the output utterances")
    parser.add_argument("--seed", type=int, default=42, help="Seed")
    parser.add_argument("--temperature", type=int, default=0.7, help="Sampling softmax temperature")
    parser.add_argument("--top_k", type=int, default=0, help="Filter top-k tokens before sampling (<=0: no filtering)")
    parser.add_argument("--top_p", type=float, default=0.9,
                        help="Nucleus filtering (top-p) before sampling (<=0.0: no filtering)")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__file__)
    logger.info(pformat(args))

    if args.model_checkpoint == "":
        logging.error("Checkpoint needed!")
        return

    random.seed(args.seed)
    torch.random.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    logger.info("Get pretrained model and tokenizer")
    tokenizer_class = BertTokenizer
    model_class = OpenAIGPTLMHeadModel if not args.gpt2 else GPT2LMHeadModel
    tokenizer = tokenizer_class.from_pretrained(args.model_checkpoint, do_lower_case=True)
    model = model_class.from_pretrained(args.model_checkpoint)

    model.to(args.device)
    model.eval()

    def tokenize(obj):
        if isinstance(obj, str):
            return tokenizer.convert_tokens_to_ids(tokenizer.tokenize(obj))
        if isinstance(obj, dict):
            return dict((n, tokenize(o)) for n, o in obj.items())
        return list(tokenize(o) for o in obj)

    history = []
    random_reply = ["hh,你接着说","这...","我想想","开心点","有点难，我之后查查相关资料"]
    blacktokens = []
    with open("blacktokens.txt","r",encoding="utf-8") as f:
        for i in f:
            blacktokens.append(i.strip())
    print(blacktokens)
               
    while True:
        raw_text = input(">>> ")
        while not raw_text:
            print('Prompt should not be empty!')
            raw_text = input(">>> ")
        raw_text = " ".join(list(raw_text.replace(" ", "")))
        history.append(tokenize(raw_text))
        with torch.no_grad():
            out_ids = sample_sequence(history, tokenizer, model, args)
        history.append(out_ids)
        history = history[-(2 * args.max_history + 1):]
        out_text = tokenizer.decode(out_ids, skip_special_tokens=True)
        out = ''.join(out_text.split())
        #print(out_text)

        out = out.replace("付费","找专人")
        
        #print(out)
        for i in blacktokens:
            if out.find(i) > 0:
                print(random.choice(random_reply))
                break
        else:
            print(out)
        # print(" ".join(jieba.cut(out)))

        # query = out_text
        #
        # # Generate salt and sign
        # salt = random.randint(32768, 65536)
        # sign = make_md5(appid + query + str(salt) + appkey)
        #
        # # Build request
        # headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        # payload = {'appid': appid, 'q': query, 'from': "zh", 'to': "en", 'salt': salt, 'sign': sign}
        #
        # # Send request
        # r = requests.post(url, params=payload, headers=headers)
        # result = r.json()
        # rst = eval(json.dumps(result, indent=4, ensure_ascii=False))["trans_result"][0]["dst"]
        # print(rst)
        #
        # args1 = run_pplm.extract_parsm()
        # convc = rst
        # args1.cond_text = convc
        # print(args1)
        # res = run_pplm.run_pplm_example(**vars(args1))
        # print(res)
        #
        # query = res[0].replace("\n","")
        # print(query)
        #
        # # Generate salt and sign
        # salt = random.randint(32768, 65536)
        # sign = make_md5(appid + query + str(salt) + appkey)
        #
        # # Build request
        # headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        # payload = {'appid': appid, 'q': query, 'from': "en", 'to': "zh", 'salt': salt, 'sign': sign}
        #
        # # Send request
        # r = requests.post(url, params=payload, headers=headers)
        # result = r.json()
        # rst = eval(json.dumps(result, indent=4, ensure_ascii=False))["trans_result"][0]["dst"]
        # print(rst)


#终端运行
if __name__ == "__main__":
    run()

# IP地址调用运行
# if __name__ == "__main__":
#     parser = ArgumentParser()
#     parser.add_argument('--gpt2', action='store_true', help="use gpt2")
#     parser.add_argument("--model_checkpoint", type=str, default="", help="Path, url or short name of the model")
#     parser.add_argument("--max_history", type=int, default=2, help="Number of previous utterances to keep in history")
#     parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
#                         help="Device (cuda or cpu)")
#
#     parser.add_argument("--no_sample", action='store_true', help="Set to use greedy decoding instead of sampling")
#     parser.add_argument("--max_length", type=int, default=30, help="Maximum length of the output utterances")
#     parser.add_argument("--min_length", type=int, default=1, help="Minimum length of the output utterances")
#     parser.add_argument("--seed", type=int, default=42, help="Seed")
#     parser.add_argument("--temperature", type=int, default=0.7, help="Sampling softmax temperature")
#     parser.add_argument("--top_k", type=int, default=0, help="Filter top-k tokens before sampling (<=0: no filtering)")
#     parser.add_argument("--top_p", type=float, default=0.9,
#                         help="Nucleus filtering (top-p) before sampling (<=0.0: no filtering)")
#     args = parser.parse_args()
#
#     logging.basicConfig(level=logging.INFO)
#     logger = logging.getLogger(__file__)
#     logger.info(pformat(args))
#
#     if args.model_checkpoint == "":
#         logging.error("Checkpoint needed!")
#         #return
#     random.seed(args.seed)
#     torch.random.manual_seed(args.seed)
#     torch.cuda.manual_seed(args.seed)
#
#     logger.info("Get pretrained model and tokenizer")
#     tokenizer_class = BertTokenizer
#     model_class = OpenAIGPTLMHeadModel if not args.gpt2 else GPT2LMHeadModel
#     tokenizer = tokenizer_class.from_pretrained(args.model_checkpoint, do_lower_case=True)
#     model = model_class.from_pretrained(args.model_checkpoint)
#
#     model.to(args.device)
#     model.eval()
#
#
#     def tokenize(obj):
#         if isinstance(obj, str):
#             return tokenizer.convert_tokens_to_ids(tokenizer.tokenize(obj))
#         if isinstance(obj, dict):
#             return dict((n, tokenize(o)) for n, o in obj.items())
#         return list(tokenize(o) for o in obj)
#
#     app.run(host='0.0.0.0', port = 8001, debug=True)